import tushare as ts
from mongoengine import *
import datetime
import time
import random
from multiprocessing import Pool


connect('DTest00')


class Stocks(Document):
    # date   open   high  close    low     volume  price_change  p_change
    # ma5    ma10    ma20      v_ma5     v_ma10     v_ma20
    date = DateTimeField()
    open = FloatField()
    high = FloatField()
    close = FloatField()
    low = FloatField()
    volume = FloatField()
    price_change = FloatField()
    p_change = FloatField()
    ma5 = FloatField()
    ma10 = FloatField()
    ma20 = FloatField()
    v_ma5 = FloatField()
    v_ma10 = FloatField()
    v_ma20 = FloatField()

    meta = {
        'ordering': ['date']
    }


def save_stock(stock_code):
    print('saving ', stock_code, '...')

    time.sleep(random.randint(1, 3))
    stock_info = ts.get_hist_data(stock_code)
    if stock_info is not None:
        for index, row in stock_info.iterrows():
            year, month, day = index.split('-')
            date = datetime.date(int(year), int(month), int(day))
            stock = Stocks.objects(date=date)
            if stock:
                stock = stock[0]
            else:
                stock = Stocks()
            stock.switch_collection(str(stock_code))
            # set data
            stock.date = date
            stock.open = float(row.open)
            stock.high = float(row.high)
            stock.close = row.close
            stock.low = row.low
            stock.volume = row.volume
            stock.price_change = row.price_change
            stock.p_change = row.p_change
            stock.ma5 = row.ma5
            stock.ma10 = row.ma10
            stock.ma20 = row.ma20
            stock.v_ma5 = row.v_ma5
            stock.v_ma10 = row.v_ma10
            stock.v_ma20 = row.v_ma20
            stock.save()
            pass
        print('complete saving ', stock_code)
    pass


def main():
    data = ts.get_stock_basics()
    stocks_list = data.index.values.tolist()
    # pool = Pool()
    # pool.map(save_stock, stocks_list)
    for stock in stocks_list:
        save_stock(stock)
    # save_stock(stocks_list[1])
    pass


if __name__ == '__main__':
    main()
